import matplotlib.pyplot as plt
import torch
from torchvision import datasets, transforms
from torch.utils.data import DataLoader, Subset

transforms = transforms.Compose([
    transforms.Resize((32, 32), antialias=True),
    transforms.ToTensor(),
])
# 读取数据集(没有就下载)
train_ds = datasets.MNIST(root='../day 1/xiazai',
                          train=True,
                          transform=transforms,
                          target_transform=lambda label: torch.tensor(label),
                          download=False)
test_ds = datasets.MNIST(root='../day 1/xiazai',
                         train=False,
                         transform=transforms,
                         target_transform=lambda label: torch.tensor(label),
                         download=False)

# 数据集加载器
train_loader = DataLoader(train_ds, batch_size=64, shuffle=True)
test_loader = DataLoader(test_ds, batch_size=128, shuffle=True)

if __name__ == '__main__':
    # 显示训练集 第一批次的前六张图片
    examples = enumerate(train_loader)
    batch_idx, (imgs, label) = next(examples)
    for i in range(6):
        plt.subplot(2, 3, i + 1)
        plt.imshow(imgs[i][0], cmap='gray')
        plt.title(label[i].item())
        plt.axis('off')
    plt.show()
    print(len(train_ds))
    print(len(train_loader))

